from backtrader import indicator
from numpy.core.fromnumeric import size
import pandas as pd
from pandas.core.indexes import period
from pymongo import MongoClient
from datetime import datetime
import time
import math
import backtrader as bt
import backtrader.indicators as btind
import backtrader.feeds as btfeeds
import talib
import numpy as np
from backtrader.feeds import GenericCSVData

objJson = {
    "_id": 'uuidstr',
    "adjustflag": 'adjustflag',
    "code": 'code',
    'start': 'start',
    'end': 'end',
    'startcash': 'startcash',
    'endtcash': 'cerebro.broker.getvalue()',
    'com': 'com',
    'qts': 'qts',
    'basicsData': [],  # 基础数据
    'isbuyArr': [],  # 买入一笔 买入价格
    'issellArr': [],  # 卖出一笔  卖出价格
    'pnlcommArr': [],  # 一笔交易后的盈亏
    'currentCapital': [],  # 资金走势
    'indicatorLines': {},  # 用户所选指标线

    'annualReturn': [],  # 年度回报
    'Calmar': [],  # 卡码率
    'DrawDown': {},  # 回撤
    'sharperatio': math.nan,  # 夏普比率
    'sqnIndex': {},  # SQN指数
    'TimeReturn': [],  # 基于时间范围内回报
}


class PandasData_more(bt.feeds.PandasData):
    lines = ('rsi',) # 要添加的线
    # 设置 line 在数据源上的列位置
    params=(
        ('rsi', 6),
           ) 
    # -1表示自动按列明匹配数据，也可以设置为线在数据源中列的位置索引 (('pe',6),('pb',7),)

# 创建策略继承bt.strategy
class TestStrategy(bt.Strategy):

    params = (
        ('maperiod', 15),
        # 判断是否输出该日志
        ('printlog', False),
    )
    

    # 初始化部分指标

    def __init__(self):

        # 指标进行注册
        self.indicators = {}
        # 跟踪挂单
        self.order = None
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None

        self.RB_h_Rsi = self.datas[2].rsi  #主期货大周期对RSI
        self.Rb_m_Rsi = self.datas[3].rsi  #主期货小周期对RSI

        self.IL_h_Rsi = self.datas[0].rsi  #对冲期货大周期对RSI
        self.Il_m_Rsi = self.datas[1].rsi  #对冲期货小周期对RSI

        self.maxRsi = 60  #最大偏离度
        self.minRsi = 10    #最小偏离度


    # 策略核心代码
    def next(self):

        self.h_Rsi = self.RB_h_Rsi[0] - self.IL_h_Rsi[0]  #大周期偏离度
        self.m_Rsi = self.Rb_m_Rsi[0] - self.Il_m_Rsi[0]  #小周期偏离度

        # print(self.m_Rsi)
        current_date = self.datas[0].datetime.date(0).isoformat()  # 当前日期
        current_capital = self.broker.getvalue()  # 当前剩余资金

        
        if (self.h_Rsi > 0 and 
            self.h_Rsi > self.minRsi and 
            self.h_Rsi < self.maxRsi and 
            self.m_Rsi > self.minRsi ):
            
            if (self.getposition(self.datas[1]).size == 0):
                self.buy(data = self.datas[1],size = 10)

            if (self.getposition(self.datas[3]).size != 0):
                self.sell(data = self.datas[3],size = 13)

            # print(self.getposition(self.datas[1]).size)
        # print(self.getposition(self.datas[1]).size)
        if (self.h_Rsi < 0 and 
           abs(self.h_Rsi) > self.minRsi and 
           abs(self.h_Rsi) < self.maxRsi and
           self.m_Rsi < 0 and
           abs(self.m_Rsi) > self.minRsi):

           if (self.getposition(self.datas[1]).size != 0):
               self.sell(data = self.datas[1],size = 10)
            
           if (self.getposition(self.datas[3]).size == 0):
               self.buy(data = self.datas[3],size = 13)

                # self.buy(data = self.datas[1],size = 10)

            # self.buy(data = self.datas[3],size = 13)
            # self.sell(data = self.datas[1],size = 10)
           
        
        

    # 日志输出

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                # print('买入价格:{}---买入手续费{}'.format(order.executed.price,
                #       order.executed.comm))
                objJson['isbuyArr'].append(
                    {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm

                print({'状态':'买入','date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
            elif order.issell():
                # print('卖出价格:{}---卖出手续费{}'.format(order.executed.price,
                #       order.executed.comm))
                objJson['issellArr'].append(
                    {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
                print({'状态':'卖出','date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消/保证金不足/拒绝', doprint=False)
        # 其他状态记录为：无法挂单
        self.order = None

    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        objJson['pnlcommArr'].append(
            {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(trade.pnlcomm, 2)})

    # 策略结束时，多用于参数调优
    def stop(self):
        # 回测结束存储数据
        # print(objJson)
        self.log('(MA均线： %2d日) 期末总资金 %.2f' %
                 (self.params.maperiod, self.broker.getvalue()), doprint=False)


def main(codeArr, startcash=1000000, com=0.00006, qts=100):

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)
    # cerebro.optstrategy(TestStrategy,maperiod=range(10, 15))
    # 获取数据

    for item in codeArr:
        if(item[1] == 'h'):
            df = pd.read_excel(r'../期货对冲相关性数据/' + item[0] + '.xls')
            df['code'] = item[0]
            df['rsi'] = talib.RSI(df.close, timeperiod=28)
            df.index = pd.to_datetime(df.date)
            # df['date2'] = df.index

            df = df[['open', 'high', 'low', 'close', 'volume','code','rsi']]
            # print(df)
            dfnew = pd.DataFrame(np.repeat(df.values,60,axis=0))
            dfnew.columns = df.columns
            dfnew['date'] = pd.DataFrame(pd.date_range(start="2010-01-01",periods=23520,freq="D"))
            dfnew.index =dfnew['date']
            # print(dfnew.head(10000))
             # 将数据加载至回测系统
            data = PandasData_more(
                dataname=dfnew.head(10000),
                # fromdate='2010-01-01',
                # todate='2037-05-18',
                # dtformat=('%Y-%m-%d'),
                # datetime=0,
                # high=2,
                # low=3,
                # open=1,
                # close=4,
                # volume=5,
                openinterest=-1
            )
            
            # dfnew.to_csv('D:/晨乐量化/backtrader_template/期货对冲相关性数据/test/' + item[0] + '.csv',sep=',',index=True,header=True)
            cerebro.adddata(data, name=item[0])
        else:
            df = pd.read_excel(r'../期货对冲相关性数据/' + item[0] + '.xls')
            df['code'] = item[0]
            df['rsi'] = talib.RSI(df.close, timeperiod=28)
            
            df['date'] = pd.DataFrame(pd.date_range(start="2010-01-01",periods=22530,freq="D"))
            df = df[['open', 'high', 'low', 'close', 'volume','code','rsi','date']]
            df.index =df['date']
            # print(df.head(10000))
            data = PandasData_more(
                dataname=df.head(10000),
                # fromdate='2010-01-01',
                # todate='2037-05-18',
                # dtformat=('%Y-%m-%d'),
                # datetime=0,
                # high=2,
                # low=3,
                # open=1,
                # close=4,
                # volume=5,
                openinterest=-1
            )

            # df.to_csv('D:/晨乐量化/backtrader_template/期货对冲相关性数据/test/' + item[0] + '.csv',sep=',',index=True,header=True)
            cerebro.adddata(data, name=item[0])


    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    cerebro.broker.set_coc(True)

    print('期初总资金: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    # strat = result[0]


    print('期末总资金: %.2f' % cerebro.broker.getvalue())
    
    cerebro.plot()


# 时间处理
def loadDate(str):
    date = datetime.strptime(str, '%Y%m%d').strftime('%Y-%m-%d')
    return (date)


def datestr2num(s):
    return datetime.strptime(s.decode('ascii'), "%d-%m-%Y").date().weekday()



def writeExcl():
    IL8_h1 = pd.read_excel(r'../期货对冲相关性数据/IL8(h1).xls')
    IL8_h1['code'] = 'IL8_h1'
    IL8_m1 = pd.read_excel(r'../期货对冲相关性数据/IL8(m1).xls')
    IL8_m1['code'] = 'IL8_m1'

    RBL_h1 = pd.read_excel(r'../期货对冲相关性数据/RBL8(h1).xls')
    RBL_h1['code'] = 'RBL_h1'
    RBL_m1 = pd.read_excel(r'../期货对冲相关性数据/RBL8(m1).xls')
    RBL_m1['code'] = 'RBL_m1'

    print(RBL_m1)




codeArr = [['IL8(h1)','h'],['IL8(m1)','m'],['RBL8(h1)','h'],['RBL8(m1)','m']]
# codeArr = [['RBL8(m1)','m'],['RBL8(h1)','h']]
main(codeArr)



